from sqlmodel import SQLModel, Field, Relationship
from datetime import datetime
from typing import Optional, List


class TaskStatus(str):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    PAUSED = "paused"


class Task(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(..., description="任务名称")
    description: Optional[str] = Field(default=None, description="任务描述")
    engine_name: str = Field(..., description="使用的模糊测试引擎")
    target_id: int = Field(..., description="目标ID")
    config: str = Field(..., description="引擎配置（JSON格式）")
    status: str = Field(default=TaskStatus.PENDING, description="任务状态")
    progress: float = Field(default=0.0, description="任务进度")
    celery_task_id: Optional[str] = Field(default=None, description="Celery任务ID")
    created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
    started_at: Optional[datetime] = Field(default=None, description="开始时间")
    completed_at: Optional[datetime] = Field(default=None, description="完成时间")
    
    # 关系
    results: List["TaskResult"] = Relationship(back_populates="task")
    artifacts: List["Artifact"] = Relationship(back_populates="task")


class TaskResult(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    task_id: int = Field(..., foreign_key="task.id")
    result_type: str = Field(..., description="结果类型（崩溃、漏洞等）")
    result_data: str = Field(..., description="结果数据（JSON格式）")
    severity: Optional[str] = Field(default=None, description="严重性")
    found_at: datetime = Field(default_factory=datetime.now, description="发现时间")
    
    # 关系
    task: Task = Relationship(back_populates="results")


class Artifact(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    task_id: int = Field(..., foreign_key="task.id")
    artifact_type: str = Field(..., description="制品类型（日志、崩溃样本等）")
    file_path: str = Field(..., description="文件路径")
    s3_key: str = Field(..., description="MinIO中的key")
    size: int = Field(..., description="文件大小")
    created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
    
    # 关系
    task: Task = Relationship(back_populates="artifacts")